import torch
import torch.nn as nn
import torch.nn.functional as F
from dm02_encoder import *

# 解码器层
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout_p):
        super(DecoderLayer, self).__init__()
        self.size = size
        # 自注意力机制对象 Q=K=V
        self.self_attn = self_attn
        # 其他注意力机制对象 Q != K = V
        self.src_attn = src_attn
        # 前馈全连接层对象
        self.feed_forward = feed_forward

        self.sublayers = clones(SublayerConnection(size, dropout_p), 3)

    def forward(self, x, memory, source_mask, target_mask):
        # x 解码器输入
        # memory 编码器输出结果
        # source_mask --> pad-mask
        # target_mask --> sentence-mask
        # 结果第一个子层连接结构
        x1 = self.sublayers[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        x2 = self.sublayers[1](x1, lambda x: self.src_attn(x, memory, memory, source_mask))
        x3 = self.sublayers[2](x2, self.feed_forward)
        return x3


class Decoder(nn.Module):
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)

        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, source_mask, target_mask):
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)

        return self.norm(x)
